# -*- coding: utf-8 -*-
import base64
import logging
import os
import sys
import traceback
import time
from email import encoders
from email.header import Header
from email.mime.base import MIMEBase
from email.mime.multipart import MIMEMultipart
from email.mime.image import MIMEImage
from email.mime.text import MIMEText
from email.utils import formataddr
from types import TracebackType
from typing import Optional
import tempfile
 
from google.auth.transport.requests import Request
from google.oauth2.credentials import Credentials
from google_auth_oauthlib.flow import InstalledAppFlow
from googleapiclient.discovery import build
from googleapiclient.errors import HttpError
import webbrowser
 
logging.basicConfig(level=logging.INFO)
 
 
class GoogleEmail:
    def __init__(self, client_secret_path: str = "", token_path: str = "", headless: bool = False) -> None:
        """
        Initialize an instance of the GoogleEmail class.
        Args:
            client_secret_path (str): The path of the client secret json file. Defaults to "".
            token_path (str): The path for saving/loading the token. If empty, uses default path. Defaults to "".
            headless (bool): Whether to use console mode for OAuth (for servers without browser). Defaults to False.
        Returns:
            None
        """
        self._client_secret_path = client_secret_path
        self._token_path = token_path
        self._headless = headless
        self._gmail_service = None
        self._init()
 
    def _init(self) -> None:
        """
        Initialize the gmail_service.
        Returns:
            None
        """
        if not os.path.exists(self._client_secret_path):
            logging.error(f"客户端密钥文件不存在: {self._client_secret_path}")
            sys.exit(1)
            
        tmp_dir = os.path.join(os.path.dirname(__file__), "tmp")
        os.makedirs(tmp_dir, exist_ok=True)
        
        # 使用用户指定的token路径或默认路径
        google_token_path = self._token_path if self._token_path else os.path.abspath(os.path.join(tmp_dir, "google_email_token.json"))
        
        credentials = None
        if os.path.exists(google_token_path):
            try:
                credentials = Credentials.from_authorized_user_file(
                    filename=google_token_path,
                    scopes=["https://www.googleapis.com/auth/gmail.send"]
                )
            except Exception as e:
                logging.error(f"读取token文件失败: {e}\n{traceback.format_exc()}")
                sys.exit(1)
                
        if not credentials or not credentials.valid:
            if credentials and credentials.expired and credentials.refresh_token:
                try:
                    credentials.refresh(Request())
                except Exception as e:
                    logging.error(f"刷新token失败: {e}\n{traceback.format_exc()}")
                    # 如果刷新失败，删除旧token文件并重新授权
                    if os.path.exists(google_token_path):
                        os.remove(google_token_path)
            else:
                try:
                    flow = InstalledAppFlow.from_client_secrets_file(
                        client_secrets_file=self._client_secret_path,
                        scopes=["https://www.googleapis.com/auth/gmail.send"]
                    )
                    
                    if self._headless:
                        # 无头模式下显示URL并接收输入的验证码
                        logging.info(f"请在浏览器中打开以下URL进行授权:")
                        logging.info(flow.authorization_url()[0])
                        code = input("请输入授权码: ")
                        flow.fetch_token(code=code)
                        credentials = flow.credentials
                    else:
                        try:
                            # 尝试检测是否有可用的浏览器
                            webbrowser.get()
                            credentials = flow.run_local_server(port=0)
                        except webbrowser.Error:
                            logging.warning("未检测到可用的浏览器，切换到手动授权模式")
                            logging.info(f"请在浏览器中打开以下URL进行授权:")
                            logging.info(flow.authorization_url()[0])
                            code = input("请输入授权码: ")
                            flow.fetch_token(code=code)
                            credentials = flow.credentials
                except Exception as e:
                    logging.error(f"OAuth授权失败: {e}\n{traceback.format_exc()}")
                    sys.exit(1)
                    
            with open(google_token_path, "w") as f:
                f.write(credentials.to_json())
 
        try:
            # 设置超时参数
            self._gmail_service = build("gmail", "v1", credentials=credentials, cache_discovery=False)
        except Exception as e:
            logging.error(f"创建Gmail服务失败: {e}\n{traceback.format_exc()}")
            sys.exit(1)
 
    def send(self,
             subject: str,
             body: str,
             to_recipients: str,
             cc_recipients: Optional[str] = None,
             bcc_recipients: Optional[str] = None,
             attachment_path: Optional[str] = None,
             sender_name: Optional[str] = None,
             sender_email: Optional[str] = None,
             retry_times: int = 3,
             retry_interval: int = 5,
             timeout: int = 30,
             inline_images: Optional[dict] = None,
             process_inline_images: bool = True  # 新增：自动处理 base64 图片
             ) -> None:
        """
        Send an email using Gmail API.
        Args:
            subject (str): The email subject.
            body (str): The email body.
            to_recipients (str): Comma-separated email addresses of the primary recipients.
            cc_recipients (str): Comma-separated email addresses of the CC recipients. Default is None.
            bcc_recipients (str): Comma-separated email addresses of the BCC recipients. Default is None.
            attachment_path (str): Path to the file to be attached. Default is None (no attachment).
            sender_name (str): Display name for the sender. Default is None (uses email address).
            sender_email (str): Email address to display with sender name. Default is None (uses "me").
            retry_times (int): Number of retries if sending fails. Default is 3.
            retry_interval (int): Interval between retries in seconds. Default is 5.
            timeout (int): Timeout for HTTP requests in seconds. Default is 30.
        Returns:
            None
        """
        # 保证所有超链接 href 有 http/https 前缀
        import re
        def fix_links(html):
            return re.sub(r'<a\s+href=["\'](?!http)(www\.[^"\'>]+)["\']', r'<a href="https://\1"', html)
        body = fix_links(body)
        
        message = MIMEMultipart('related')  # ✅ 关键：使用 'related' 类型
        message.preamble = 'This is a multi-part message in MIME format.'

        # ✅ 正确设置 Subject：直接赋值 Header 对象，不要 str()
        message["Subject"] = Header(subject, 'utf-8')

        message["To"] = to_recipients
        if cc_recipients:
            message["Cc"] = cc_recipients
        if bcc_recipients:
            message["Bcc"] = bcc_recipients
        if sender_name:
            email = sender_email if sender_email else "me"
            message["From"] = formataddr((sender_name, email))

        # --- 新增：处理 base64 图片 ---
        extracted_inline_images = {}
        temp_dir = None

        if process_inline_images:
            # 创建临时目录保存提取的图片
            temp_dir = tempfile.mkdtemp(prefix="email_images_")
            img_counter = 0
            def replace_base64(match):
                if not match:
                    logging.warning("正则匹配为空，跳过")
                    return match.group(0) if match else ''

                full_match = match.group(0)  # 完整匹配的 src="..."
                # 提取引号内的内容
                # 支持 "..." 或 '...'
                if full_match.startswith("src='"):
                    uri = full_match[5:-1]  # 去掉 src=' 和 '
                elif full_match.startswith("src=\""):
                    uri = full_match[5:-1]  # 去掉 src=" 和 "
                else:
                    uri = full_match[4:]  # src=... 无引号？尽量提取
                    uri = uri.strip().strip("'\"")

                # 解析 
                header_match = re.match(r'data:([^;]+)(?:;[^,]+)?,?(.*)', uri, re.DOTALL)
                if not header_match:
                    logging.warning(f"无法解析 Data URI: {uri[:50]}...")
                    return full_match

                mime_type = header_match.group(1)  # image/jpeg
                base64_data = header_match.group(2) or ''

                # 推断扩展名
                ext = mime_type.split('/')[-1]
                if ext not in ['png', 'jpg', 'jpeg', 'gif', 'bmp', 'webp']:
                    ext = 'png'

                # 解码 base64
                try:
                    image_data = base64.b64decode(base64_data, validate=True)
                except Exception as e:
                    logging.error(f"Base64 解码失败: {e}")
                    return full_match

                nonlocal img_counter, extracted_inline_images
                cid = f"img{img_counter}"
                img_counter += 1

                img_path = os.path.join(temp_dir, f"{cid}.{ext}")
                try:
                    with open(img_path, "wb") as f:
                        f.write(image_data)
                    extracted_inline_images[cid] = img_path
                    return f'src="cid:{cid}"'  # 返回完整的 src="cid:xxx"
                except Exception as e:
                    logging.error(f"保存临时图片失败: {e}")
                    return full_match

            # ✅ 使用改进的正则表达式
            body = re.sub(
                r'src=["\']\s*data:image/[^"\';]+(?:;[^"\']*)?,[^"\']*?["\']',
                replace_base64,
                body,
                flags=re.IGNORECASE | re.DOTALL
                )
            logging.info(f"自动提取并转换了 {len(extracted_inline_images)} 个 base64 图片为 cid 引用")
        # --- end ---

        # 构建正文（包含内联图片引用）
        html_part = MIMEText(body, "html", "utf-8")
        message.attach(html_part)

        # 添加用户手动指定的内联图片
        final_inline_images = {}
        if inline_images:
            final_inline_images.update(inline_images)
        # 添加自动提取的图片
        if extracted_inline_images:
            final_inline_images.update(extracted_inline_images)

        # 添加所有内联图片作为附件
        if final_inline_images:
            for cid, img_path in final_inline_images.items():
                if not os.path.exists(img_path):
                    logging.warning(f"内联图片不存在: {img_path}")
                    continue
                with open(img_path, "rb") as f:
                    img = MIMEImage(f.read())
                img.add_header('Content-ID', f'<{cid}>')
                img.add_header('Content-Disposition', 'inline')
                message.attach(img)
                logging.info(f"内联图片已添加: cid:{cid} -> {img_path}")

        # 添加普通附件（如果有）
        if attachment_path:
            if not os.path.exists(attachment_path):
                logging.error(f"附件不存在: {attachment_path}")
                return 0
            attachment_name = os.path.basename(attachment_path)
            with open(attachment_path, "rb") as f:
                attachment = MIMEBase("application", "octet-stream")
                attachment.set_payload(f.read())
            encoders.encode_base64(attachment)
            attachment.add_header(
                "Content-Disposition",
                f"attachment; filename={Header(attachment_name, 'utf-8').encode()}"
            )
            message.attach(attachment)


        # ✅ 生成 raw message：使用 as_bytes() 并指定 encoding
        try:
            raw_message = base64.urlsafe_b64encode(message.as_bytes()).decode("utf-8")
        except Exception as e:
            logging.error(f"序列化邮件失败: {e}")
            return 0
    
        if self._gmail_service is None:
            logging.error("Gmail服务未初始化")
            return 0
        
        for attempt in range(retry_times):
            try:
                logging.info(f"尝试发送邮件 (尝试 {attempt+1}/{retry_times})...")
                request = self._gmail_service.users().messages().send(
                    userId="me", 
                    body={"raw": raw_message}
                )
                # 设置请求超时
                request.http.timeout = timeout
                response = request.execute()
                logging.info(f"邮件发送成功, 响应: {response}")
                return 1
            except HttpError as e:
                logging.error(f"发送邮件时出现HTTP错误: {e.status_code} - {e.reason}")
                if e.status_code == 429 or e.status_code >= 500:  # 限流或服务器错误
                    if attempt < retry_times - 1:
                        logging.info(f"将在 {retry_interval} 秒后重试...")
                        time.sleep(retry_interval)
                    else:
                        logging.error("达到最大重试次数，放弃发送")
                else:
                    logging.error(f"发送邮件失败: {e}")
                    break
            except Exception as e:
                logging.error(f"发送邮件失败: {e}\n{traceback.format_exc()}")
                if attempt < retry_times - 1:
                    logging.info(f"将在 {retry_interval} 秒后重试...")
                    time.sleep(retry_interval)
                else:
                    logging.error("达到最大重试次数，放弃发送")

        return 0
 
    @staticmethod
    def generate_token(client_secret_path: str, token_output_path: str) -> None:
        """
        在本地生成令牌文件，用于上传到无浏览器环境的服务器
        Args:
            client_secret_path (str): 客户端密钥文件路径
            token_output_path (str): 令牌输出路径
        Returns:
            None
        """
        if not os.path.exists(client_secret_path):
            logging.error(f"客户端密钥文件不存在: {client_secret_path}")
            sys.exit(1)
            
        try:
            flow = InstalledAppFlow.from_client_secrets_file(
                client_secrets_file=client_secret_path,
                scopes=["https://www.googleapis.com/auth/gmail.send"]
            )
            credentials = flow.run_local_server(port=0)
            with open(token_output_path, "w") as f:
                f.write(credentials.to_json())
            logging.info(f"令牌已成功生成到: {token_output_path}")
        except Exception as e:
            logging.error(f"生成令牌失败: {e}\n{traceback.format_exc()}")
            sys.exit(1)
 
 
if __name__ == "__main__":
    import argparse
    
    parser = argparse.ArgumentParser(description="Google Email Sender")
    parser.add_argument("--generate-token", action="store_true", help="仅生成令牌文件")
    parser.add_argument("--client-secret", type=str, default="/tmp/client_secret.json", help="客户端密钥文件路径")
    parser.add_argument("--token-path", type=str, default="", help="令牌文件路径")
    parser.add_argument("--headless", action="store_true", help="使用无头模式（无浏览器环境）")
    parser.add_argument("--to", type=str, default="", help="收件人邮箱")
    parser.add_argument("--subject", type=str, default="Test Email", help="邮件主题")
    parser.add_argument("--body", type=str, default="This is a test email.", help="邮件内容")
    parser.add_argument("--sender-name", type=str, default="", help="发件人显示名称")
    parser.add_argument("--sender-email", type=str, default="", help="发件人邮箱地址")
    parser.add_argument("--timeout", type=int, default=30, help="请求超时时间（秒）")
    
    args = parser.parse_args()
    
    if args.generate_token:
        token_path = args.token_path if args.token_path else os.path.join(os.path.dirname(__file__), "tmp", "google_email_token.json")
        GoogleEmail.generate_token(args.client_secret, token_path)
    else:
        if not args.to:
            logging.error("请指定收件人邮箱地址")
            sys.exit(1)
            
        # 发送邮件
        GoogleEmail(
            client_secret_path=args.client_secret,
            token_path=args.token_path,
            headless=args.headless
        ).send(
            to_recipients=args.to,
            subject=args.subject,
            body=args.body,
            sender_name=args.sender_name,
            sender_email=args.sender_email,
            timeout=args.timeout
        )